import os.path

import torch

from ClassifyModel import ClassifyModel


def save_state(model: ClassifyModel, current_epoch: int):
    torch.save(model.custom_state_dict(), f"model_state_checkpoint{current_epoch}.pth")
    print(f"Save model state to model_state_checkpoint{current_epoch}.pth")


# Deprecated
def save_all(model, current_epoch: int):
    torch.save(model, f"model_all_checkpoint{current_epoch}.pth")
    print(f"Save model to model_all_checkpoint{current_epoch}.pth")
    save_state(model, current_epoch)


def load_from_state(model: ClassifyModel, epoch: int):
    model.load_custom_dict(torch.load(f"model_state_checkpoint{epoch}.pth"))
    print(f"Load model state from model_state_checkpoint{epoch}.pth")


# Deprecated
def load_all(epoch: int):
    print(f"Loading model from model_all_checkpoint{epoch}.pth ...")
    return torch.load(f"model_all_checkpoint{epoch}.pth")


def has_checkpoint(epoch: int) -> bool:
    return os.path.exists(f"model_state_checkpoint{epoch}.pth") or os.path.exists(f"model_all_checkpoint{epoch}.pth")
